import os
import torch
import torch.nn as nn
import torch.optim as optim
import scipy.io
import matplotlib.pyplot as plt
import numpy as np
from torch.utils.data import Dataset, DataLoader, random_split

# 选择设备
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# -------------------- 参数设置 --------------------
# m=1 表示每个样本仅一帧数据，多批次训练时 batch_size > 1
start = 1                       # 文件起始编号
num_samples = 2               # 假设总共有 100 个样本
file_prefix = "spatial_correlation_"
file_extension = ".mat"
data_folder = "dync"
k = 481                       # 丰度矩阵维度（输出稀疏表示维度）
num_epochs = 1000               # 训练轮数
batch_size = 10                 # 每个批次训练样本数
train_ratio = 0.8               # 80% 作为训练集，其余作为测试集

# 端元矩阵文件路径
D_path = "Data/DC1/phasecha.mat"

# -------------------- 自定义数据集 --------------------
class MatDataset(Dataset):
    def __init__(self, folder, prefix, extension, start_idx, num_samples):
        """
        加载文件编号从 start_idx 开始，总共 num_samples 个样本，
        每个样本为 (2,8,8) 的双通道数据（分别为实部和虚部）。
        """
        self.folder = folder
        self.prefix = prefix
        self.extension = extension
        self.start_idx = start_idx
        self.num_samples = num_samples
        self.samples = []
        for i in range(start_idx, start_idx + num_samples):
            fname = os.path.join(folder, f"{prefix}{i}{extension}")
            if os.path.exists(fname):
                mat = scipy.io.loadmat(fname)
                if "p" in mat:
                    Y = mat["p"]
                    if Y.size == 64:
                        # reshape 成 (8,8) 后分别取实部与虚部，组成 (2,8,8)
                        img = np.stack([np.real(Y.reshape(8,8)), np.imag(Y.reshape(8,8))], axis=0)
                        self.samples.append(img)
                    else:
                        print(f"Shape of 'p' in {fname} is not 64.")
                else:
                    print(f"'p' not found in {fname}")
            else:
                print(f"File not found: {fname}")
        if len(self.samples) == 0:
            raise ValueError("No samples found. Check file indices and paths.")
    
    def __len__(self):
        return len(self.samples)
    
    def __getitem__(self, idx):
        sample = self.samples[idx]
        return torch.from_numpy(sample).float()  # (2,8,8)

# -------------------- 数据加载与划分 --------------------
dataset = MatDataset(data_folder, file_prefix, file_extension, start, num_samples)
train_size = int(train_ratio * len(dataset))
test_size = len(dataset) - train_size
train_dataset, test_dataset = random_split(dataset, [train_size, test_size])
train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)

def load_endmember_matrix(file_path):
    """
    加载端元矩阵，并将其转换为形状为 (128, k) 的张量，
    其中 128 = 64 实部 + 64 虚部。
    """
    mat = scipy.io.loadmat(file_path)
    D = mat["phasecha"]
    if D.shape[0] == 64:
        D_reshaped = D.reshape(64, k)
    else:
        raise ValueError("The shape of D is not compatible with 8x8 reshaping.")
    D_real = np.real(D_reshaped)
    D_imag = np.imag(D_reshaped)
    D_combined = np.concatenate((D_real, D_imag), axis=0)  # (128, k)
    return torch.from_numpy(D_combined).float().to(device)

D_tensor = load_endmember_matrix(D_path)

# -------------------- 网络结构 --------------------
# 针对 m=1 的情况，每个样本只有一帧数据，采用纯CNN提取空间特征

class SingleStepNet(nn.Module):
    def __init__(self, in_channels, cnn_out_channels, k):
        """
        输入为 (B,2,8,8) 的批次数据，
        使用 CNN 提取空间特征，然后全连接层映射到稀疏表示 X，形状为 (B,k)。
        """
        super(SingleStepNet, self).__init__()
        self.cnn = SpatialCNNBranch(in_channels, cnn_out_channels)
        self.fc = nn.Linear(cnn_out_channels, k)
        self.relu = nn.ReLU()
        self.dropout = nn.Dropout(0.1)
        
    def forward(self, x):
        # x: (B,2,8,8)
        feat = self.cnn(x)       # (B, cnn_out_channels)
        X = self.fc(feat)        # (B, k)
        X = self.relu(X)
        X = self.dropout(X)
        # 稀疏化：对每个样本输出中保留前 k 个最大值（假设 fc 输出维度 >= k）
        topk_values, topk_indices = torch.topk(X, k, dim=1)
        mask = torch.zeros_like(X)
        mask.scatter_(1, topk_indices, topk_values)
        X = mask / (torch.sum(mask, dim=1, keepdim=True) + 1e-8)
        return X

class SpatialCNNBranch(nn.Module):
    def __init__(self, in_channels, out_channels):
        super(SpatialCNNBranch, self).__init__()
        self.conv1 = nn.Conv2d(in_channels, 64, kernel_size=3, padding=1)
        self.bn1   = nn.BatchNorm2d(64)
        self.relu  = nn.ReLU()
        self.pool  = nn.MaxPool2d(2,2)
        self.conv2 = nn.Conv2d(64, 128, kernel_size=3, padding=1)
        self.bn2   = nn.BatchNorm2d(128)
        self.pool2 = nn.MaxPool2d(2,2)
        # 利用 dummy 数据计算全连接层输入维度
        dummy_input = torch.zeros(1, in_channels, 8, 8)
        dummy_output = self.pool2(self.relu(self.bn2(self.conv2(self.pool(self.relu(self.bn1(self.conv1(dummy_input))))))))
        feature_dim = dummy_output.view(1, -1).size(1)
        self.fc = nn.Linear(feature_dim, out_channels)
        
    def forward(self, x):
        x = self.conv1(x)
        x = self.bn1(x)
        x = self.relu(x)
        x = self.pool(x)
        x = self.conv2(x)
        x = self.bn2(x)
        x = self.relu(x)
        x = self.pool2(x)
        x = x.view(x.size(0), -1)
        x = self.fc(x)
        return x

# -------------------- 模型初始化 --------------------
in_channels = 2         # 双通道输入
cnn_out_channels = 128  # CNN 输出维度
model = SingleStepNet(in_channels, cnn_out_channels, k).to(device)
criterion = nn.MSELoss()
optimizer = optim.Adam(model.parameters(), lr=5e-5, weight_decay=1e-5)

# -------------------- 训练过程 --------------------
loss_history = []
for epoch in range(num_epochs):
    model.train()
    epoch_loss = 0.0
    for batch_data in train_loader:
        batch_data = batch_data.to(device)  # (B, 2, 8, 8)
        optimizer.zero_grad()
        X_train = model(batch_data)          # (B, k)
        # 重构过程：X_train 的转置 (k, B) 与端元矩阵 D (128, k) 矩阵乘得到 (128, B)
        X_T = X_train.transpose(0, 1)        # (k, B)
        Y_reconstructed = torch.matmul(D_tensor, X_T)  # (128, B)
        Y_reconstructed = Y_reconstructed.transpose(0, 1)  # (B, 128)
        # 将原始 batch 数据展平成 (B, 128) ，因为 2*8*8 = 128
        Y_true = batch_data.view(batch_data.size(0), -1)  # (B, 128)
        loss = criterion(Y_reconstructed, Y_true)
        loss.backward()
        optimizer.step()
        epoch_loss += loss.item() * batch_data.size(0)
    epoch_loss /= len(train_dataset)
    loss_history.append(epoch_loss)
    if (epoch + 1) % 100 == 0:
        print(f"Epoch [{epoch+1}/{num_epochs}], Loss: {epoch_loss:.6f}")

plt.figure()
plt.plot(loss_history)
plt.xlabel("Epoch")
plt.ylabel("Loss")
plt.title("Training Loss Curve")
plt.show()

# -------------------- 测试集评估 --------------------
model.eval()
with torch.no_grad():
    total_loss = 0.0
    all_rmse = []
    for batch_data in test_loader:
        batch_data = batch_data.to(device)
        X_eval = model(batch_data)           # (B, k)
        X_T_eval = X_eval.transpose(0, 1)      # (k, B)
        Y_reconstructed = torch.matmul(D_tensor, X_T_eval)  # (128, B)
        Y_reconstructed = Y_reconstructed.transpose(0, 1)   # (B, 128)
        Y_true = batch_data.view(batch_data.size(0), -1)      # (B, 128)
        loss_eval = criterion(Y_reconstructed, Y_true)
        total_loss += loss_eval.item() * batch_data.size(0)
        rmse = torch.sqrt(torch.mean((Y_reconstructed - Y_true) ** 2, dim=1))
        all_rmse.append(rmse.cpu())
    total_loss /= len(test_dataset)
    all_rmse = torch.cat(all_rmse, dim=0)
    print(f"Final Reconstruction MSE (Test): {total_loss:.6f}")
    print("RMSE per sample:", all_rmse)
    
# 绘制测试集中第一个批次第一个样本的重构与真实结果对比
with torch.no_grad():
    sample = next(iter(test_loader)).to(device)  # 取一个测试批次
    X_sample = model(sample)
    X_T_sample = X_sample.transpose(0, 1)
    Y_reconstructed_sample = torch.matmul(D_tensor, X_T_sample).transpose(0, 1)
    Y_true_sample = sample.view(sample.size(0), -1)
    
plt.figure()
plt.plot(Y_reconstructed_sample[0].cpu().numpy(), label="Reconstructed")
plt.plot(Y_true_sample[0].cpu().numpy(), label="Ground Truth")
plt.legend()
plt.title("Comparison of one test sample (flattened)")
plt.show()


